import sys
import heapq
from dataclasses import dataclass, field
from enum import IntEnum, Enum
from typing import List

import pygame
from loguru import logger


class E_MapType(IntEnum):
    """
    地图标记物
    """

    Blank: int = 0
    Wall: int = 1
    StartPoint: int = 2
    EndPoint: int = 4
    APath: int = 8


class E_MapColor(Enum):
    """
    地图不同标记物颜色, 颜色顺序RGB
    """

    Gap = (255, 255, 255)  # 间隔-白色
    Blank = (0, 255, 0)  # 空白-绿色
    Wall = (0, 0, 0)  # 墙-黑色
    StartPoint = (0, 0, 255)  # 起点-蓝色
    EndPoint = (255, 0, 0)  # 终点-红色
    APath = (255, 0, 255)  # A*路径-紫色


@dataclass
class MapNode:
    """
    地图点-属性
    """

    x: int = 0  # x轴坐标
    y: int = 0  # y轴坐标
    map_type: int = 0  # 地图类型 0:空白 1:墙 2:起点 3:终点
    f: int = 0  # 总距离
    g: int = 0  # 起点到当前点的距离
    h: int = 0  # 当前点到终点的距离
    parent: any = None  # 父节点

    def __lt__(self, other):
        return self.f < other.f

    def __str__(self):
        return f"MapNode(x={self.x} y={self.y} type={self.map_type} f={self.f} g={self.g} h={self.h})"


@dataclass
class MapData:
    """
    地图数据
    """

    w: int = 10  # 宽
    h: int = 10  # 高
    screenScale: int = 50  # 屏幕缩放比例
    lineSize: int = 3  # 网格线宽
    data: List[MapNode] = field(default_factory=list)


class gameRect:
    def __init__(self, mapData: MapData):
        pygame.init()
        self.map = mapData
        self.screenScale = self.map.screenScale
        self.width = self.map.w * self.screenScale
        self.height = self.map.h * self.screenScale
        self.screen = pygame.display.set_mode((self.width, self.height))
        self.clock = pygame.time.Clock()
        self.init_map()
        self.init_font()
        self.init_data()

    def __exit__(self):
        pygame.quit()

    def init_map(self):
        """
        初始化地图数据
        """
        self.map.data = []
        # 遍历行 * 列
        for x in range(self.map.w):
            map_list = []
            for y in range(self.map.h):
                map = MapNode()
                map.x = x
                map.y = y
                map.map_type = E_MapType.Blank.value
                map.total_dist = 0
                map.g_dist = 0
                map.h_dist = 0
                map_list.append(map)
            self.map.data.append(map_list)
        print(self.map.data)

    def reset_map(self):
        # 遍历行 * 列
        for x in range(self.map.w):
            for y in range(self.map.h):
                map = self.map.data[x][y]
                map.map_type = E_MapType.Blank.value
                map.total_dist = 0
                map.g_dist = 0
                map.h_dist = 0
                map.parent = None

    def init_font(self):
        """初始化字体"""
        font_path = "/usr/share/fonts/truetype/wqy/wqy-microhei.ttc"
        font_size = 16
        self.font = pygame.font.Font(font_path, font_size)

    def init_data(self):
        self.mapType = E_MapType.Wall
        self.mouse_pos = (0, 0)  # 鼠标屏幕位置
        self.startPoint = None  # 起点 地图
        self.endPoint = None  # 终点 地图

    def draw_map(self):
        """
        绘制网格线
        """
        lineSize = self.map.lineSize
        self.startPoint = None
        self.endPoint = None
        # 遍历行 * 列
        for i in range(self.map.w):
            for j in range(self.map.h):
                # 计算当前方格左上角坐标
                x = i * self.screenScale + lineSize
                y = j * self.screenScale + lineSize
                map_type = self.map.data[i][j].map_type
                if map_type == E_MapType.Blank.value:
                    color = E_MapColor.Blank.value
                elif map_type == E_MapType.Wall.value:
                    color = E_MapColor.Wall.value
                elif map_type == E_MapType.StartPoint.value:
                    color = E_MapColor.StartPoint.value
                    self.startPoint = self.map.data[i][j]
                elif map_type == E_MapType.EndPoint.value:
                    color = E_MapColor.EndPoint.value
                    self.endPoint = self.map.data[i][j]
                elif map_type == E_MapType.APath.value:
                    color = E_MapColor.APath.value
                    self.endPoint = self.map.data[i][j]
                else:
                    continue
                pygame.draw.rect(
                    self.screen, color, (x, y, self.screenScale - lineSize, self.screenScale - lineSize), 0
                )

        # 绘制启点
        if self.startPoint:
            x = self.startPoint.x * self.screenScale
            y = self.startPoint.y * self.screenScale
            self.draw_text("start", (x, y + self.screenScale // 3), (255, 255, 255))
        # 绘制终点
        if self.endPoint:
            x = self.endPoint.x * self.screenScale
            y = self.endPoint.y * self.screenScale
            self.draw_text("end", (x, y + self.screenScale // 3), (255, 255, 255))

    def draw_note_text(self):
        """
        绘制提示文字: wall,start,end
        """
        pos = self.mouse_pos
        color = (255, 255, 255)
        if self.mapType == E_MapType.Wall:
            text = "wall"
        elif self.mapType == E_MapType.StartPoint:
            text = "start"
        elif self.mapType == E_MapType.EndPoint:
            text = "end"
        self.draw_text(text, pos, color)

    def draw_text(self, text, pos, color):
        text_surface = self.font.render(text, True, color)
        self.screen.blit(text_surface, pos)

    def get_map_idx_by_screen_pos(self, pos):
        """
        根据屏幕坐标获取地图坐标
        """
        x = pos[0] // self.screenScale
        y = pos[1] // self.screenScale
        return x, y

    def get_map_by_screen_pos(self, pos):
        """
        根据屏幕坐标获取地图数据
        """
        x = pos[0] // self.screenScale
        y = pos[1] // self.screenScale
        return self.map.data[x][y]

    def handler_mouse_event(self, event):
        """
        鼠标按下事件
            左键: 标记wall, start, end
            右键: 清空标记
        """
        # NOTE: numpy 结构化数据类型, 索引返回引用,可以直接修改
        # map_idx = self.get_map_by_screen_pos(event.pos)
        # if event.button == 1:  # 左键
        #     map_type = self.mapType.value
        #     self.map.data[map_idx[0]][map_idx[1]]["map_type"] = map_type
        # elif event.button == 3:
        #     self.map.data[map_idx[0]][map_idx[1]]["map_type"] = E_MapType.Blank.value

        map = self.get_map_by_screen_pos(event.pos)
        if event.button == 1:  # 左键
            map_type = self.mapType.value
            map.map_type = map_type
        elif event.button == 3:
            map.map_type = E_MapType.Blank.value

    def handler_mouse_motion(self, event):
        """
        鼠标移动事件
        """
        self.mouse_pos = event.pos

    def handler_key_press(self, event):
        """
        按键按下
            1 标记类型wall
            2 标记类型start
            3 标记类型end
            q 求相邻点
            w 计算A星路径
            c 清空地图
        """
        if event.key == pygame.K_1:
            self.mapType = E_MapType.Wall
        elif event.key == pygame.K_2:
            self.mapType = E_MapType.StartPoint
        elif event.key == pygame.K_3:
            self.mapType = E_MapType.EndPoint
        elif event.key == pygame.K_q:
            cur_map = self.get_map_by_screen_pos(self.mouse_pos)
            near_map = self.get_near_map(cur_map)
            print(near_map)
        elif event.key == pygame.K_w:
            path_list = self.a_start_search(self.startPoint, self.endPoint)
            print("A星路径:")
            for path in path_list:
                print(path)
                path.map_type = E_MapType.APath.value
        elif event.key == pygame.K_c:
            self.reset_map()

    def loop_run(self):
        self.color_idx = 0
        while True:
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    pygame.quit()
                    sys.exit()
                elif event.type == pygame.MOUSEBUTTONDOWN:
                    self.handler_mouse_event(event)
                elif event.type == pygame.MOUSEMOTION:
                    self.handler_mouse_motion(event)
                elif event.type == pygame.KEYDOWN:
                    self.handler_key_press(event)

            color = E_MapColor.Gap.value
            self.screen.fill(color)  # 清屏
            self.draw_map()  # 绘制地图
            self.draw_note_text()

            pygame.display.update()  # 刷新屏幕
            self.clock.tick(1000)  # 限制fps 1帧

    ### ---------------------------------------
    # 地图算法实现
    ### ---------------------------------------
    def get_near_map(self, map):
        """
        获取当前地图的上下左右四个邻居
        """
        cur_x = map.x
        cur_y = map.y
        near_map = []
        y = cur_y
        # 左
        x = cur_x - 1
        if x >= 0:
            if self.map.data[x][y].map_type != E_MapType.Wall.value:
                near_map.append(self.map.data[x][y])
        # 右
        x = cur_x + 1
        if x < self.map.w:
            if self.map.data[x][y].map_type != E_MapType.Wall.value:
                near_map.append(self.map.data[x][y])

        x = cur_x
        # 上
        y = cur_y - 1
        if y >= 0:
            if self.map.data[x][y].map_type != E_MapType.Wall.value:
                near_map.append(self.map.data[x][y])
        # 下
        y = cur_y + 1
        if y < self.map.h:
            if self.map.data[x][y].map_type != E_MapType.Wall.value:
                near_map.append(self.map.data[x][y])

        return near_map

    def cal_map_distance(self, start: MapNode, end: MapNode):
        """
        返回: 曼哈顿距离
        """
        return abs(start.x - end.x) + abs(start.y - end.y)

    def a_start_search(self, start, end):
        if (not start or start.map_type != E_MapType.StartPoint.value) or (
            not end or end.map_type != E_MapType.EndPoint.value
        ):
            logger.error("start or end is not valid")
            return None

        open_list = []
        close_set = set()
        heapq.heappush(open_list, (0, start))
        while open_list:
            _, cur_map = heapq.heappop(open_list)
            if cur_map == end:
                a_path = []
                while cur_map:
                    a_path.append(cur_map)
                    cur_map = cur_map.parent
                return a_path

            close_set.add((cur_map.x, cur_map.y))
            for near in self.get_near_map(cur_map):
                if (near.x, near.y) in close_set or near.map_type == E_MapType.Wall.value:
                    continue
                near.g = cur_map.g + 1
                near.h = self.cal_map_distance(near, end)
                near.f = near.g + near.h
                near.parent = cur_map
                heapq.heappush(open_list, (near.f, near))


if __name__ == "__main__":
    map_data = MapData()
    gameImpl = gameRect(map_data)
    gameImpl.loop_run()
